#pragma once

#include <iostream>
#include <stdexcept>
#include <vector>

// #include "TxnnCpp/Common.h"
#include "utils/data/image_data.h"

namespace det {

/**
 * @brief 掩码结构体, 存储掩码数据及其尺寸信息
 */
struct Mask {
  std::vector<uint8_t> data; // 掩码数据
  int width = 0;             // 掩码宽度
  int height = 0;            // 掩码高度

  Mask()                                        = default; // 默认构造函数
  Mask(const Mask& other)                       = default; // 拷贝构造函数
  Mask& operator=(const Mask& other)            = default; // 拷贝赋值运算符
  Mask(Mask&& other) noexcept                   = default; // 移动构造函数
  Mask& operator=(Mask&& other) noexcept        = default; // 移动赋值运算符

  Mask(int width, int height) : width(width), height(height) {
    if (width < 0 || height < 0) {
      throw std::invalid_argument(
        ERROR_MSG("Mask width and height must be positive"));
    }
    data.resize(width * height);
  }

  friend std::ostream& operator<<(std::ostream& os, const Mask& mask) {
    os << "Mask(width=" << mask.width << ", height=" << mask.height << ", data size=" << mask.data.size() << ")";
    return os;
  }
};

/**
 * @brief 矩形框结构体, 存储矩形框的坐标信息
 */
struct Box {
  float left    = 0.0f; // 左边坐标
  float top     = 0.0f; // 上边坐标
  float right   = 0.0f; // 右边坐标
  float bottom  = 0.0f; // 下边坐标

  Box()                                       = default; // 默认构造函数
  Box(const Box& other)                       = default; // 拷贝构造函数
  Box& operator=(const Box& other)            = default; // 拷贝赋值运算符
  Box(Box&& other) noexcept                   = default; // 移动构造函数
  Box& operator=(Box&& other) noexcept        = default; // 移动赋值运算符

  Box(float left, float top, float right, float bottom)
    : left(left), top(top), right(right), bottom(bottom) {}

  friend std::ostream& operator<<(std::ostream& os, const Box& box) {
    os << "Box(left=" << box.left << ", top=" << box.top << ", right=" << box.right << ", bottom=" << box.bottom << ")";
    return os;
  }
};

/**
 * @brief 基础结果结构体, 存储检测结果的基础信息, 如数量、类别和得分
 */
struct BaseResult {
  int count = 0;                // 检测到的目标数量
  std::vector<int> classes;     // 检测到的目标类别
  std::vector<float> scores;    // 检测到的目标得分

  BaseResult()                                        = default; // 默认构造函数
  BaseResult(const BaseResult& other)                 = default; // 拷贝构造函数
  BaseResult& operator=(const BaseResult& other)      = default; // 拷贝赋值运算符
  BaseResult(BaseResult&& other) noexcept             = default; // 移动构造函数
  BaseResult& operator=(BaseResult&& other) noexcept  = default; // 移动赋值运算符

  BaseResult(int count, const std::vector<int>& classes, const std::vector<float>& scores)
    : count(count), classes(classes), scores(scores) {}
};

/**
 * @brief 分类结果结构体, 继承自基础结果结构体, 无需额外构造函数
 */
struct ClassificationResult : public BaseResult {
  friend std::ostream& operator<<(std::ostream& os, const ClassificationResult& result) {
    os << "ClassificationResult(\n    num=" << result.count << ",\n    classes=[";
    for (const auto& c : result.classes) { os << c << ", "; }
    os << "],\n    scores=[";
    for (const auto& s : result.scores) { os << s << ", "; }
    os << "])";
    return os;
  }
};

/**
 * @brief 检测结果结构体, 继承自基础结果结构体, 增加矩形框信息
 */
struct DetectionResult : public BaseResult {
  std::vector<Box> boxes; // 检测到的目标矩形框

  DetectionResult()                                                 = default; // 默认构造函数
  DetectionResult(const DetectionResult& other)                     = default; // 拷贝构造函数
  DetectionResult& operator=(const DetectionResult& other)          = default; // 拷贝赋值运算符
  DetectionResult(DetectionResult&& other) noexcept                 = default; // 移动构造函数
  DetectionResult& operator=(DetectionResult&& other) noexcept      = default; // 移动赋值运算符

  DetectionResult(int count, const std::vector<int>& classes, const std::vector<float>& scores, const std::vector<Box>& boxes)
    : BaseResult(count, classes, scores), boxes(boxes) {}

  friend std::ostream& operator<<(std::ostream& os, const DetectionResult& result) {
    os << "DetectionResult(\n    num=" << result.count << ",\n    classes=[";
    for (const auto& c : result.classes) { os << c << ", "; }
    os << "],\n    scores=[";
    for (const auto& s : result.scores) { os << s << ", "; }
    os << "],\n    boxes=[\n";
    for (const auto& b : result.boxes) { os << "        " << b << ",\n"; }
    os << "    ]\n)";
    return os;
  }
};

} // namespace det